package com.websocket.service;

import com.websocket.WebSocketConfig;
import com.websocket.event.EventPublisher;
import com.websocket.exception.service.ExceptionDispatch;
import com.websocket.handshake.HandshakeExecutionChain;
import com.websocket.service.dispatch.WebSocketDispatch;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.java_websocket.WebSocket;
import org.java_websocket.drafts.Draft;
import org.java_websocket.exceptions.InvalidDataException;
import org.java_websocket.framing.CloseFrame;
import org.java_websocket.handshake.ClientHandshake;
import org.java_websocket.handshake.ServerHandshakeBuilder;

import java.net.InetSocketAddress;
import java.security.Principal;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class WebSocketServer extends org.java_websocket.server.WebSocketServer {

    private static final Log logger = LogFactory.getLog(WebSocketServer.class);

    private final WebSocketConfig webSocketConfig;
    private final EventPublisher eventPublisher;
    private final WebSocketDispatch websocketDispatch;
    private final ExceptionDispatch exceptionDispatch;
    private static final Map<Principal, WebSocket> connections = new ConcurrentHashMap<>();

    public WebSocketServer(WebSocketConfig webSocketConfig, WebSocketDispatch websocketDispatch, ExceptionDispatch exceptionDispatch) {
        super(new InetSocketAddress(webSocketConfig.handshakeRegistry.getPort()));
        this.webSocketConfig = webSocketConfig;
        this.websocketDispatch = websocketDispatch;
        this.exceptionDispatch = exceptionDispatch;
        this.eventPublisher = webSocketConfig.eventPublisher;
    }

    public WebSocket getConnection(Principal user) {
        return connections.get(user);
    }

    @Override
    public ServerHandshakeBuilder onWebsocketHandshakeReceivedAsServer(WebSocket socket, Draft draft, ClientHandshake handshake) throws InvalidDataException {
        try {
            String endpoint = webSocketConfig.handshakeRegistry.getEndpoint();
            if (endpoint != null) {
                String reqEndpoint = handshake.getResourceDescriptor();
                if (!reqEndpoint.startsWith(endpoint))
                    throw new InvalidDataException(CloseFrame.TLS_ERROR);
            }
            Principal user = webSocketConfig.determineUser.apply(handshake);
            socket.setAttachment(new Attachment(user));
            // 握手拦截器链
            HandshakeExecutionChain interceptorChain = webSocketConfig.handshakeRegistry.getInterceptorChain();
            if (!interceptorChain.applyHandle(new WebSocketConnection(socket), draft, handshake, connections)) {
                throw new InvalidDataException(CloseFrame.TLS_ERROR);
            }
            return super.onWebsocketHandshakeReceivedAsServer(socket, draft, handshake);
        } catch (Exception ex) {
            exceptionDispatch.doDispatch(ex, null);
        }
        throw new InvalidDataException(CloseFrame.TLS_ERROR);
    }

    @Override
    public void onOpen(WebSocket conn, ClientHandshake handshake) {
        Attachment attachment = conn.getAttachment();
        Principal user = attachment.getPrincipal();
        connections.put(user, conn);

        // 发布事件
        eventPublisher.publishConnectEvent(this, user);
    }

    @Override
    public void onClose(WebSocket conn, int code, String reason, boolean remote) {
        Attachment attachment = conn.getAttachment();
        Principal user = attachment.getPrincipal();
        connections.remove(user);
        // 发布事件
        eventPublisher.publishDisconnectEvent(this, user, code);
    }

    @Override
    public void onMessage(WebSocket conn, String message) {
        Attachment attachment = conn.getAttachment();
        Principal user = attachment.getPrincipal();
        websocketDispatch.doDispatch(conn, message, user);
        // 发布事件
        eventPublisher.publishReceiveEvent(this, user, message);
    }

    @Override
    public void onError(WebSocket conn, Exception ex) {
        exceptionDispatch.doDispatch(ex, conn);
    }

    @Override
    public void onStart() {
        logger.info("Websocket initialized with port: " + this.getPort());
    }

}
